﻿using LxBasic;

namespace LxLESS
{
    public class LESSSolver
    {
        public IterationTypes IterateType { get; set; } = IterationTypes.AMGV;
        public int PreSmoothSteps { get; set; } = 7;
        public int PostSmoothSteps { get; set; } = 1;
        private int maxCoarseLevel = -1;
        private static LESSSolver _LessSolver = new LESSSolver();
        public static void Solve<T1, T2>(IEnumerable<T1> equations, IEnumerable<T2> connects, IterationTypes iteType = IterationTypes.AMGV, int preSmoothSteps = 5, int postSmoothStep = 3)
          where T1 : IZoneIteration where T2 : IJointIteration
        {
            _LessSolver.IterateType = iteType;
            _LessSolver.PreSmoothSteps = preSmoothSteps;
            _LessSolver.PostSmoothSteps = postSmoothStep;
            _LessSolver.Solve(equations, connects);
        }
        public void Solve<T1, T2>(IEnumerable<T1> equations, IEnumerable<T2> connects) where T1 : IZoneIteration where T2 : IJointIteration
        {
            AMGGlobal.SolvingLevel = -1;
            switch (IterateType)
            {
                case IterationTypes.AMGV:
                    iniAMG();
                    AMG_V(equations, connects);//V循环
                    break;
                case IterationTypes.AMGF:
                    iniAMG();
                    AMG_F(equations, connects);//F循环
                    break;
                case IterationTypes.AMGW:
                    iniAMG();
                    AMG_W(equations, connects);//W循环
                    break;
                case IterationTypes.Jacobi:
                    SolveByJacobi(equations, connects, PreSmoothSteps + PostSmoothSteps);
                    break;
                case IterationTypes.BiCGSTAB:
                    SolveByBiCG(equations, connects, PreSmoothSteps + PostSmoothSteps);
                    break;
                case IterationTypes.Explicit:
                    SolveExplicit(equations, connects, PreSmoothSteps + PostSmoothSteps);
                    break;
            }
            void iniAMG()
            {
                Array.Fill(AMGGlobal.IsNeededUpdateA, true);
                maxCoarseLevel = -1;
                foreach (var e in equations)
                {
                    var lessZone = e.GetLessZone();
                    lessZone.IsSymmetric = e.IsSymmetric();
                    int lev = lessZone.MaxCoarseLevel;
                    if (lev > maxCoarseLevel)
                    {
                        maxCoarseLevel = lev;
                    }
                }
            }
        }

        public void SolveExplicit<T1, T2>(IEnumerable<T1> equations, IEnumerable<T2> connects, int cycles = 1) where T1 : IZoneExplicitIteration where T2 : IJointIteration
        {
            MultiThreads.RunParallelTask((threadInfo) =>
            {
                foreach (var c in connects) { c.Iterate_Copy_xArray(threadInfo); }
                foreach (var s in equations) { s.ExplicitIterate_Step1(threadInfo); }
            });
            MultiThreads.RunParallelTask((threadInfo) => { foreach (var s in equations) { s.ExplicitIterate_Step2(threadInfo); } });
        }
        private void Smooth<T1, T2>(IEnumerable<T1> equations, IEnumerable<T2> connects, int cycles = 1) where T1 : IZoneIteration where T2 : IJointIteration
        {
            if (AMGGlobal.SolvingLevel == -1)
            {
                SolveByJacobi(equations, connects);
            }
            else
            {
                if (cycles <= 2)
                {
                    SolveByJacobi(equations, connects, cycles);
                }
                else
                {
                    SolveByBiCG(equations, connects, cycles);
                }
            }
        }
        private void SolveByJacobi<T1, T2>(IEnumerable<T1> equations, IEnumerable<T2> connects, int cycles = 1) where T1 : IZoneIteration where T2 : IJointIteration
        {
            cycles = (cycles + 1) / 2;//jacobi一个循环实际有2次迭代
            for (int i1 = 0; i1 < cycles; i1++)
            {
                LoopConnects(connects, (c) => c.Iterate_Copy_xArray);
                LoopEquations(equations, (z) => z.JacobiIterate_Step1);
                LoopConnects(connects, (c) => c.Iterate_Copy_xCommonArray);
                LoopEquations(equations, (z) => z.JacobiIterate_Step2);
            }
        }
        public void CalResidual<T1, T2>(IEnumerable<T1> equations, IEnumerable<T2> connects) where T1 : IZoneIteration where T2 : IJointIteration
        {
            LoopConnects(connects, (c) => c.Iterate_Copy_xArray);
            LoopEquations(equations, (z) => z.CalResidual);
        }
        private void SolveByBiCG<T1, T2>(IEnumerable<T1> equations, IEnumerable<T2> connects, int cycles = 5) where T1 : IZoneIteration where T2 : IJointIteration
        {
            LoopConnects(connects, (c) => c.Iterate_Copy_xArray);
            LoopEquations(equations, (z) => z.BiCGStabIte_Ini_r_rho1);
            if (BiCGStabGlobal.Initialize_Step2())
            {
                SolveByJacobi(equations, connects);
            }
            else
            {
                LoopEquations(equations, (z) => z.BiCGStabIte_Ini_r0_pArray_yArray);
                for (int i1 = 1; ; i1++)
                {
                    LoopConnects(connects, (c) => c.BiCGStabIte_Copy_yArray);
                    LoopEquations(equations, (z) => z.BiCGStabIte_vArray_alpha1);
                    if (BiCGStabGlobal.Cal_alpha()) return;
                    LoopEquations(equations, (z) => z.BiCGStabIte_sArray_zArray);
                    LoopConnects(connects, (c) => c.BiCGStabIte_Copy_zArray);
                    LoopEquations(equations, (z) => z.BiCGStabIte_tArray_omega1);
                    if (BiCGStabGlobal.Cal_omega()) return;
                    LoopEquations(equations, (z) => z.BiCGStabIte_xArray);
                    if (i1 >= cycles || BiCGStabGlobal.GetIsNeedBreak()) { return; }
                    LoopEquations(equations, (z) => z.BiCGStabIte_rArray_rho1);
                    BiCGStabGlobal.Cal_rhobeta();
                    LoopEquations(equations, (z) => z.BiCGStabIte_pArray_yArray);
                }
            }
        }
        private void AMG_V<T1, T2>(IEnumerable<T1> equations, IEnumerable<T2> connects) where T1 : IZoneIteration where T2 : IJointIteration
        {
            int AMG_V_StartLevel = AMGGlobal.SolvingLevel;
            while (AMGGlobal.SolvingLevel < maxCoarseLevel)
            {
                Smooth(equations, connects, PreSmoothSteps);
                RestrictAndLevelUp(equations, connects);
            }
            Smooth(equations, connects, PreSmoothSteps + PostSmoothSteps);
            while (AMGGlobal.SolvingLevel > AMG_V_StartLevel)
            {
                LevelDownAndReceiveProlong(equations);
                Smooth(equations, connects, PostSmoothSteps);
            }
        }
        private void AMG_F<T1, T2>(IEnumerable<T1> equations, IEnumerable<T2> connects) where T1 : IZoneIteration where T2 : IJointIteration
        {
            Smooth(equations, connects);//-1级上smooth
            while (AMGGlobal.SolvingLevel < maxCoarseLevel)
            {
                RestrictAndLevelUp(equations, connects);
            }
            Smooth(equations, connects, PreSmoothSteps + PostSmoothSteps);//max级上smooth
            LevelDownAndReceiveProlong(equations);
            while (AMGGlobal.SolvingLevel >= 0)
            {
                AMG_V(equations, connects);
                LevelDownAndReceiveProlong(equations);
            }
            Smooth(equations, connects);//-1级上smooth
        }
        private void AMG_W<T1, T2>(IEnumerable<T1> equations, IEnumerable<T2> connects) where T1 : IZoneIteration where T2 : IJointIteration
        {
            Smooth(equations, connects);//-1级上smooth
            RestrictAndLevelUp(equations, connects);
            while (AMGGlobal.SolvingLevel < maxCoarseLevel)
            {
                AMG_V(equations, connects);
                RestrictAndLevelUp(equations, connects);
            }
            Smooth(equations, connects, PreSmoothSteps + PostSmoothSteps);//max级上smooth
            LevelDownAndReceiveProlong(equations);
            while (AMGGlobal.SolvingLevel >= 0)
            {
                AMG_V(equations, connects);
                LevelDownAndReceiveProlong(equations);
            }
            Smooth(equations, connects);//-1级上smooth
        }
        private void RestrictAndLevelUp<T1, T2>(IEnumerable<T1> equations, IEnumerable<T2> connects) where T1 : IZoneIteration where T2 : IJointIteration
        {
            LoopConnects(connects, c => c.Iterate_Copy_xArray);
            LoopEquations(equations,
                z => AMGGlobal.IsNeededUpdateA[AMGGlobal.SolvingLevel + 1]
                ? z.AMGIte_Restrict_Are : z.AMGIte_Restrict_re);//当前level传递系数到下一level，所以询问下一级是否需要更新系数
            AMGGlobal.SolvingLevel++;
            AMGGlobal.IsNeededUpdateA[AMGGlobal.SolvingLevel] = false;
        }
        private void LevelDownAndReceiveProlong<T1>(IEnumerable<T1> equations) where T1 : IZoneIteration
        {
            AMGGlobal.SolvingLevel--;
            LoopEquations(equations, (z) => z.AMGIte_ReceiveProlong);
        }

        private void LoopEquations<T>(IEnumerable<T> equations, Func<IZoneIteration, Action<ThreadInfo>> getAct) where T : IZoneIteration
        {
            if (AMGGlobal.SolvingLevel == -1)
            {
                MultiThreads.RunParallelTask((threadInfo) =>
                {
                    foreach (var z in equations)
                        getAct(z).Invoke(threadInfo);
                });
            }
            else
            {
                MultiThreads.RunParallelTask((threadInfo) =>
                {
                    foreach (var z in equations)
                        getAct(z.GetLessZone().MGZones[AMGGlobal.SolvingLevel]).Invoke(threadInfo);
                });
            }
        }
        private void LoopConnects<T>(IEnumerable<T> connects, Func<IJointIteration, Action<ThreadInfo>> getAct) where T : IJointIteration
        {
            if (connects.Count() == 0) return;
            if (AMGGlobal.SolvingLevel == -1)
            {
                MultiThreads.RunParallelTask((threadInfo) =>
                {
                    foreach (var z in connects)
                        getAct(z).Invoke(threadInfo);
                });
            }
            else
            {
                MultiThreads.RunParallelTask((threadInfo) =>
                {
                    foreach (var z in connects)
                        getAct(z.GetLessJoint().MGJoints[AMGGlobal.SolvingLevel]).Invoke(threadInfo);
                });
            }
        }
    }
}